#!/usr/bin/env bash

# License:
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU Lesser General Public
#   License as published by the Free Software Foundation; either
#   version 2.1 of the License, or (at your option) any later
#   version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Lesser General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public
#   License along with this program.  If not, see
#   <https://www.gnu.org/licenses/>.
#
## You can redistribute it as you wish : GPL v3
## author : wildtruc@noneltd.net

## this script have to be launch with user's session started except for rescue 'default' option command.

# check if root and exit if not
if [[ $EUID -ne 0 ]]; then
  printf 'Script must be executed as root or sudo :\n'
## using polkit to set admin rights.
  exit 1
fi

install_dir='/etc/nvidia-prime'
nvidia32_lib_file='/etc/ld.so.conf.d/nvidia-lib32.conf'
nvidia64_lib_file='/etc/ld.so.conf.d/nvidia-lib64.conf'
conf=$1
## Default Fedora propriatary driver dirs.
nv_drv_32='/usr/lib/nvidia'
nv_drv_64='/usr/lib64/nvidia'
nv_xorg_path='/usr/lib64/nvidia/xorg'
## check any custom dirs for nvidia librairies linking
if [[ -s $install_dir/library.conf ]]; then
	## TODO Here special distro library config could take place.
	## no clue without contribs.
	## test library conf content before loading.
	if [ $(cat $install_dir/library.conf| sed -En "s/^(.*)='(.*)'$/\2/g;p"| grep -c .) -gt 0 ]; then
		. $install_dir/library.conf
	fi
fi
## default options (have to be manual set/reset in options.conf of nvidia-prime-ui)
delay=5
intel_nodelay=1
cfg_grub=1
cfg_modeset=1
nv_modeset=0

if [[ -s $install_dir/options.conf ]]; then
	. $install_dir/options.conf
fi
current_display=$(xrandr --current| grep "connected"| grep "[0-9]\{3,4\}[x]"| sed -n "s/\ (.*$//p")
connected_display=( "$(xrandr --current| grep -w "connected"| sed -En "s|^(.*)\ (\(.*)$|\1|p")" )

grub_config(){
## Since fedora 30, grub config has changed and regex is no longer adapted to new version that only
## write a 'set_default' for all kernel version options. This need to be control and for other distros aswell.
	if [ $cfg_grub -gt 0 ]; then
		grub_cfg=$(find /boot -wholename "*/grub.cfg")
		grub_def=/etc/default/grub
		grub_key='GRUB_CMDLINE_LINUX'
		blacklist='nouveau.modeset=0 rd.driver.blacklist=nouveau'
		nov_set=$(cat $grub_def| egrep -c "^$grub_key.*blacklist=nouveau.*$")
		drm_set=$(cat $grub_def| egrep -c "^$grub_key.*nvidia-drm.modeset.*$")
		## if 'nvidia-drm.modeset' is define in options conf, do
		if [ $cfg_modeset -gt 0 ]; then
			## if 'blacklist=nouveau' if already defined in grub defautl conf, just add drm set.
			## esle insert the whole pattern.
			if [ $nov_set -gt 0 ]; then
				modset='nvidia-drm.modeset=0'
			else
				modset='nvidia-drm.modeset=0 '$blacklist
			fi
		else
			modset=$blacklist
		fi
		## define if grub s already setted. Then, act as appropriate. 
		cfg_prev=$(cat $grub_def| egrep -c "^$grub_key.*$modeset.*$")
		if [ $cfg_prev -eq 0 ]; then
			if [ $nov_set -eq 0 ]||[[ $cfg_modeset -gt 0 && $drm_set -eq 0 ]]; then
				## add nouveau blacklist to modeprobe.d configs list.
				if ! [[ -s /etc/modprobe.d/blacklist-nouveau.conf ]]; then
					echo "blacklist nouveau" > /etc/modprobe.d/blacklist-nouveau.conf
				fi
				sed -Eni "s|^($grub_key)=\"(.*)\"$|\1=\"$modset \2\"|g;p" $grub_def
				grub2-mkconfig -o $grub_cfg
				reboot_text='Grub has been modified for Nvidia Prime modesetting.\nPLease reboot the system for changes to take effect.'
			fi
		fi
	fi
}
bus_id_query(){
	## detect and config pci ids.
	bus_vga_list=$(lspci | egrep "\ (VGA|3D)\ "| awk '{print $1" "$5}'| sed -En "s|^.([0-9])\:.([0-9])\.([0-9])\ (.*)$|bus_\L\4=\1:\2:\3|p")
	export $bus_vga_list
}
clean_files() {
  files=(
    '/etc/X11/xorg.conf.d/00-avoid-glamor.conf'
    '/etc/X11/xorg.conf.d/99-nvidia.conf'
    '/etc/X11/xorg.conf'
  )
	## remove/backup unuseful files, but check if they exist first
	for file in "${files[@]}"; do
		if [[ -f $file ]]; then
			printf "Backing up: %s as %s\n" "$file" "$file.prime.bak"
			if [[ -s $file ]]; then cp "$file" "$file.prime.bak"; fi
			if [[ -s $file ]]; then rm -f "$file"; fi
		fi
	done
}
xrandr_outputs(){
	ifs=$IFS
	IFS=$(echo -en "\n\b")
	## intel settings only need empty xinitrc and display file. Then, only nvidia is managed.
	if [ $ctrl_key -gt 0 ]; then
		for _dsp in ${connected_display[@]}; do
			echo "$_dsp"
			output_class=$(printf "$_dsp"|awk '{print $1}'| sed -En "s|^(\w*[A-Z])[-]?(.*)$|\1|p")
			output_num=$(printf "$_dsp"|awk '{print $1}'| sed -En "s|^(\w*[A-Z])[-]?(.*)$|\2|p")
## XXX this part seems useless. Should be removed in near future if no issue reported.
#			dash_cnt=$( dash=${_dsp//[^-]}; printf "${#dash}"; )
#			## nvidia xrandr monitor set is : monitor-x-x, intel is : monitor-x
#			## counting number of dash detect the provider to config.
#			if [ $dash_cnt -eq 1 ]; then
#				output_num=$output_num'-'$output_num
#			fi
			connected_output=$output_class'-'$output_num
#			connected_output=$output_class'-'$output_num'-'$output_num
			if [ $(xrandr --current| grep "$_dsp"| egrep -oc "[0-9]{3,4}[x]") -gt 0 ]; then
				if [[ $_dsp =~ ^VGA* ]]; then
					current_output=$(xrandr --current| grep "$_dsp"| sed -En "s|^(.*)\ (\(.*)$|\1|p")
				else
					current_output=$(xrandr --current| grep "$_dsp"| sed -En "s|^(.*[0-9]{3,4}[x].*)\ (\(.*)$|\1|p")
				fi
				check_primary_set=( $(echo -e $current_output| tr ' ' '\n') )
				## command line results is ususally made of 3 word if not primary (4 otherwise)
				if [ ${#check_primary_set[@]} -gt 3 ]; then rank=4; prim=' --primary'; else rank=3; prim=''; fi
				if ! [[ $connected_output =~ ^VGA ]]; then
					def_res=$(printf "$current_output"| cut -d' ' -f$rank)
					def_horiz=$(printf "$def_res"| sed -En "s|^(.*)x(.*)\+.*\+.*$|\1|p")
					def_verti=$(printf "$def_res"| sed -En "s|^(.*)x(.*)\+.*\+.*$|\2|p")
					def_pos=$(printf "$def_res"| sed -En "s|^.*x.*\+(.*\+.*)$|\1|;s|\+|x|p")
					## don't set display as primary at this step. this here basic xrandr display setting.
					## user have to set default primary at user session in case of multiple screen.
					enable_output='--output '$connected_output' --mode '$def_horiz'x'$def_verti' --pos '$def_pos$prim
					_outputs+=($enable_output)
					if [[ $prim != '' ]]; then prim_display=$connected_output; fi
				fi
			fi
		done
		if [ $(xrandr --current| egrep -c "VGA.*[ ]connected") -gt 0 ]; then
			if [ $(echo -e "${_outputs[@]}"| grep -c "VGA") -eq 0 ]; then
				disable_output='--output VGA-0 --off'
				_outputs+=($disable_output)
				_disabled+=($disable_output)
			fi
		fi
		IFS=$ifs
		xrandr_config="xrandr ${_outputs[*]}"
		xrandr_source='NVIDIA-0'
		xrandr_sink='modesetting'
		xrandr_init_src="xrandr --setprovideroutputsource $xrandr_sink $xrandr_source"
#		xrandr_init_dsp="xrandr --auto ${_disabled[*]}"
		xrandr_init_dsp="xrandr --auto"
		## script to parse in xinitrc.prime and display.prime
		bash_start='if [ $(xrandr --current| egrep -c "'$prim_display' connected") -gt 0 ]; then'
#		bash_end='else\nxrandr --auto '${_rescue[@]}'\nfi\n'
		bash_end='fi\n'
		echo -e "#! /bin/bash\n\n#sleep $delay\n$xrandr_init_src\n$bash_start\n$xrandr_init_dsp\n$bash_end" > $install_dir/xinitrc.prime
		echo -e "#! /bin/bash\n\nsleep $delay\n$bash_start\n$xrandr_config\n$bash_end" > $install_dir/display.prime
	else
		echo -e "" > $install_dir/xinitrc.prime
		echo -e "" > $install_dir/display.prime
	fi
	## TODO ? look in monitors.xml file if exist, then edit and parse ?
	## XXX TEST offload system for render on nvidia. driver v435.17 and up compatible hardware is mandatory. 
	
	if [ $(stat -c %a $install_dir/display.prime) -ne 755 ]; then
		chmod 755 $install_dir/display.prime
	fi
	if [ $(stat -c %a $install_dir/xinitrc.prime) -ne 755 ]; then
		chmod 755 $install_dir/xinitrc.prime
	fi
}
session_manager(){
	home=/home/$(ls -l "$(pwd)"| cut -d' ' -f3 | sed -n "2p")	
	## configure each session manager if deteed: gdm kdm sddm lightdm
	session_list=(
	"lightdm,lightdm.conf,/etc/lightdm,.config/autostart,\[Seat:\*\],type.*local"
	"sddm,Xsetup,/usr/share/sddm/scripts,.config/autostart"
	"gdm,nvidia-session.desktop,/usr/share/gdm/greeter/autostart,.config/autostart"
	"kdm,Xsetup,/usr/share/config/kdm,.kde/share/autostart"
	"xdg,nvidia-session.desktop,/etc/xdg/autostart"
	)
	## loop on each session manager and configure as appropriate
	for _session in ${session_list[@]}; do
		session=$(printf "$_session"| cut -d"," -f1)
		session_conf=$(printf "$_session"| cut -d"," -f2)
		session_dir=$(printf "$_session"| cut -d"," -f3)
		session_local=$(printf "$_session"| cut -d"," -f4)
		conf_key=$(printf "$_session"| cut -d"," -f5)
		line_key=$(printf "$_session"| cut -d"," -f6)
		if [[ -d $session_dir ]]; then
			case $session in
				'lightdm')
					if ! [[ -s $session_dir/$session_conf.default ]]; then
						mv -f $session_dir/$session_conf $session_dir/$session_conf.default
						ln -sf $session_dir/$session_conf.default $session_dir/$session_conf
					fi
					case $conf in
						'nvidia')
							if [ $(cat $session_dir/$session_conf| grep -c "xinitrc.prime") -eq 0 ]; then
								if ! [[ -s $session_dir/$session_conf.prime ]]; then
									cp -f $session_dir/$session_conf.default $session_dir/$session_conf.prime
								fi
								if [ $(cat $session_dir/$session_conf.prime| grep -c "xinitrc.prime") -eq 0 ]; then
									if [ $(cat $session_dir/$session_conf.prime| grep -c "$conf_key") -gt 0 ]; then
										if [ $(cat $session_dir/$session_conf.prime| grep -c "$line_key") -gt 0 ]; then
											conf_line=$line_key
											local_conf="display-setup-script=$install_dir/xinitrc.prime"
										else
											conf_line=$conf_key
											local_conf="type=xlocal\ndisplay-setup-script=$install_dir/xinitrc.prime"
										fi
										sed -Eni "s|^[#]?($conf_line.*)$|\1\n$local_conf|g;p" $session_dir/$session_conf.prime
									else
										echo -e "[Seat:*]\ntype=xlocal\ndisplay-setup-script=$install_dir/xinitrc.prime"	>> $session_dir/$session_conf.prime
									fi
								fi
							fi
							ln -sf $session_dir/$session_conf.prime $session_dir/$session_conf
						;;
						'intel')
							ln -sf $session_dir/$session_conf.default $session_dir/$session_conf
						;;
					esac
				;;
				'gdm'|'xdg')
					case $conf in
					'nvidia')
						if ! [[ -h $session_dir/$session_conf ]]; then
							ln -sf $install_dir/$session_conf $session_dir/$session_conf
						fi
					;;
					'intel')
						if [[ -h $session_dir/$session_conf ]]; then
							rm -f $session_dir/$session_conf
						fi
					;;
					esac
				;;
				'kdm'|'sddm')
					case $conf in
					'nvidia')
						if ! [[ -h $session_dir/$session_conf ]]; then
							ln -sf $install_dir/xinitrc.prime $session_dir/$session_conf
						fi
					;;
					'intel')
						if ! [[ -h $session_dir/$session_conf ]]; then
							rm -f $session_dir/$session_conf
						fi
					;;
					esac
				;;
			esac
		fi
		if [[ -d $home/$session_local ]]; then
			if ! [[ -s $home/$session_local/nvidia-prime.desktop ]]; then
				cp -f $install_dir/nvidia-prime.desktop $home/$session_local
			fi
		fi
	done
}
abi_crtl(){
## enable xorg ABI23 behaviour to fix vsync and render. 
	if [[ $abi_ver && $abi_ver -ge 23 ]]; then
		if [ $(cat $install_dir/xorg.$conf.conf| grep -c "IgnoreABI" ) -eq 0 ]; then
			server_flag="Section \"ServerFlags\"\n\tOption\t\"IgnoreABI\" \"1\"\nEndSection"
			sed -Eni -e "s|^(## ABI config)$|\1\n$server_flag|g;p" $install_dir/xorg.$conf.conf
		fi
	fi
}
display_config(){
	xorg_file=$install_dir/xorg.$conf.conf
	if [[ $current_display ]]; then
		xrandr_outputs
		session_manager
	fi
	for vga_id in "nvidia" "intel"; do
		eval bus_id='$'bus_$vga_id
		if [ $(cat $xorg_file|grep -oc "bus_$vga_id") -gt 0 ]; then
			sed -ni -e "s|bus_$vga_id|$bus_id|g;p" $xorg_file
		else
			if [ $(cat $xorg_file| sed -n "/## $vga_id xorg/,/PCI/p"| grep -c "$bus_id") -eq 0 ]; then
				line_up=$(sed -En "/## $vga_id xorg/,/PCI/{;/BusID/p}" $xorg_file)
				sed -Eni "s/$line_up/\tBusID\t\t\"PCI:$bus_id\"/g;p" $xorg_file
			fi
		fi
	done
	ln -sf $xorg_file /etc/X11/xorg.conf
}
go_for_nvidia(){
	clean_files
#	delay=5
	ctrl_key=1
	display_config
	abi_crtl
	if [[ ! -s $nvidia64_lib_file || \
	$(cat $nvidia64_lib_file| egrep -c "$nv_drv_64") -eq 0 ]]; then
		printf "$nv_drv_64" > $nvidia64_lib_file
		if [[ -d $nv_drv_32 ]]; then
			printf "$nv_drv_32" > $nvidia32_lib_file
		fi
		ldconfig
	fi
}
go_for_intel(){
	clean_files
	if [ $intel_nodelay -eq 1 ]; then delay=0; else delay=5; fi
	ctrl_key=0
	display_config
}
go_default(){
	go_for_intel
	printf '/usr/lib64' > $nvidia64_lib_file
	if [[ -f $nvidia32_lib_file ]]; then
		printf '/usr/lib' > $nvidia32_lib_file
	fi
	ldconfig
}
go_offload(){
	echo "In test, be patient..."; exit 0
	## offload basicaly need nvida config with a different "Serverlayout" than default and glarmor loaded.
	##		Section "ServerLayout"
	##		Identifier	"layout"
	##		Screen	0	"intel_screen"
	##		Option		"AllowNVIDIAGPUScreens"
	##		EndSection
	## There is also an issue with GLX libs, nvidia's dev were introduce libglxserver_nvidia.so in place of libglx.so in ../xorg/modules/extensions with new drivers,
	## This probably solve the libglx linking for custom driver installation. Anyways, My optimus hardware support ends with 390.129 drivers that only take care of libglx.so.
	## All I can do is theorical.
	## The solution reside probably with libglxserver.so installation in default xorg location and remove the "Section files" from xorg.conf. Can't tell at this point.
	## I'll starting tests with zenvidia and see what's happening.
}
## configure grub for modesetting and blacklist all nouveau drivers.
grub_config
bus_id_query
# keeping this for debug purpose only.
# current=$(file -p /etc/X11/xorg.conf| sed -En "s|^.*xorg.(.*).conf$|\1|p")
if [[ -s /var/log/Xorg.0.log ]]; then
	abi_ver=$(cat /var/log/Xorg.0.log| sed -En "s|^.*ABI.*Driver.* ([0-9]{2})\..*$|\1|p"| uniq)
fi
case $1 in
	default ) go_default ;;
	nvidia ) go_for_nvidia ;;
	offload ) go_offload ;;
	intel ) go_for_intel ;;
#	test) xrandr_outputs; exit 0;;
	*|help )
		printf "Usage:\t$(basename $0 ) intel|nvidia|default\n\n"
		printf "\tintel\t\tSwitch to Intel GPU.\n"
		printf "\tnvidia\t\tSwitch to Nvidia GPU.\n"
		printf "\tdefault\t\tRescue to default throught commandline in case of XServer crash.\n"
		printf "\thelp\t\tthis help.\n"
		printf "\n\tYou need to logout from your session after each function.\n\n"
		exit 0
	;;
esac

#ldconfig
if [ $(file -p /etc/X11/xorg.conf| grep -c "xorg.$conf.conf") -gt 0 ]; then
	printf "\nSuccessfully switched to %s completed.\n" "$1"
	printf "Please logout for changes to take effect.\n"
else
	printf "\nSwitch to %s doesn't ended as expected\n" "$1"
	printf "Please, run script again or check custom values.\n"
fi	
